/*
 *
 * Copyright (c) 2022 The Raptor Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include "src/linux/tcp_client.h"

#include <string.h>

#include "raptor-lite/impl/connector.h"
#include "raptor-lite/impl/config.h"
#include "raptor-lite/impl/endpoint.h"
#include "raptor-lite/utils/log.h"
#include "raptor-lite/utils/time.h"

#include "src/common/endpoint_impl.h"
#include "src/common/socket_util.h"
#include "src/linux/socket_settings.h"

namespace raptor {

struct InternalClientEventNode {
    MultiProducerSingleConsumerQueue::Node node;
    NetEvent event;
    int err_code;
    Slice data;
};

#define close_socket(fd)                                                                           \
    if (fd != -1) {                                                                                \
        ::close(fd);                                                                               \
        fd = -1;                                                                                   \
    }

constexpr int TEMP_BUFFER_SIZE = 4 * 1024 * 1024;
constexpr int MIN_TIME_SLICE = 300;

TcpClient::TcpClient(ClientHandler *handler)
    : service_(handler)
    , shutdown_(true)
    , is_connected_(false)
    , fd_(-1)
    , tcp_user_timeout_ms_(3000) {

    buffer_ = new char[TEMP_BUFFER_SIZE];

    next_package_size_ = 0;
    last_sendtime_ = 0;
    last_recvtime_ = 0;

    count_.Store(0);
}

TcpClient::~TcpClient() {
    if (buffer_) delete[] buffer_;
}

raptor_error TcpClient::Init() {
    if (!shutdown_) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient is already running");
    }

    raptor_error e = epoll_.create();
    if (e != RAPTOR_ERROR_NONE) {
        return e;
    }

    bool success = false;
    mpscq_thd_ =
        Thread("client_thread", std::bind(&TcpClient::WorkThread, this, std::placeholders::_1),
               nullptr, &success);

    if (!success) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient failed to create mpscq thread");
    }

    epoll_thd_ =
        Thread("epoll_thread", std::bind(&TcpClient::EpollProcess, this, std::placeholders::_1),
               nullptr, &success);

    if (!success) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient failed to create epoll thread");
    }

    shutdown_ = false;

    return RAPTOR_ERROR_NONE;
}

raptor_error TcpClient::Start() {
    if (shutdown_) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient is not initialized");
    }

    count_.Store(0);
    next_package_size_ = 0;
    last_sendtime_ = GetCurrentMilliseconds();
    last_recvtime_ = last_sendtime_;
    mpscq_thd_.Start();
    epoll_thd_.Start();
    return RAPTOR_ERROR_NONE;
}

void TcpClient::Shutdown() {
    if (!shutdown_) {
        log_warn("TcpClient: prepare to shutdown");
        shutdown_ = true;

        epoll_.remove(fd_, EPOLLIN | EPOLLOUT | EPOLLET | EPOLLRDHUP);

        close_socket(fd_);
        is_connected_ = false;

        cv_.Signal();
        epoll_thd_.Join();
        mpscq_thd_.Join();
        epoll_.shutdown();

        // clear message queue
        bool empty = true;
        do {
            auto n = mpscq_.PopAndCheckEnd(&empty);
            if (n != nullptr) {
                delete reinterpret_cast<InternalClientEventNode *>(n);
            }
        } while (!empty);

        rcv_buffer_.ClearBuffer();
        snd_buffer_.ClearBuffer();
        log_warn("TcpClient: shutdown");
    }
}

raptor_error TcpClient::Connect(const std::string &addr) {
    raptor_resolved_addresses *addrs;
    auto e = raptor_blocking_resolve_address(addr.c_str(), nullptr, &addrs);
    if (e != RAPTOR_ERROR_NONE) {
        return e;
    }
    if (addrs->naddrs == 0) {
        return RAPTOR_ERROR_FROM_FORMAT("Invalid address: (%s)", addr.c_str());
    }

    for (int i = 0; i < static_cast<int>(addrs->naddrs); i++) {
        e = AsyncConnect(&addrs->addrs[i], tcp_user_timeout_ms_);
        if (e != RAPTOR_ERROR_NONE) {
            log_warn("TcpClient: Failed to connect %s, %s", addr.c_str(), e->ToString().c_str());
        } else {
            break;
        }
    }
    raptor_resolved_addresses_destroy(addrs);
    return e;
}

raptor_error TcpClient::AsyncConnect(const raptor_resolved_address *addr, int tcp_user_timeout_ms) {
    raptor_resolved_address mapped_addr;
    int sock_fd = -1;

    raptor_error result =
        raptor_tcp_client_prepare_socket(addr, &mapped_addr, &sock_fd, tcp_user_timeout_ms);
    if (result != RAPTOR_ERROR_NONE) {
        return result;
    }
    int err = 0;
    do {
        err = connect(sock_fd, (const raptor_sockaddr *)mapped_addr.addr, mapped_addr.len);
    } while (err < 0 && errno == EINTR);

    if (errno != EWOULDBLOCK && errno != EINPROGRESS) {
        raptor_set_socket_shutdown(sock_fd);
        return RAPTOR_POSIX_ERROR("connect");
    }

    fd_ = sock_fd;

    epoll_.add(sock_fd, reinterpret_cast<void *>(sock_fd),
               EPOLLIN | EPOLLOUT | EPOLLET | EPOLLRDHUP);

    char *str_addr = nullptr;
    raptor_sockaddr_to_string(&str_addr, &mapped_addr, 0);
    if (str_addr) {
        log_info("TcpClient: start connecting %s", str_addr);
        free(str_addr);
    }
    return RAPTOR_ERROR_NONE;
}

void TcpClient::EpollProcess(void *) {
    while (!shutdown_) {
        int count = epoll_.polling(MIN_TIME_SLICE);

        for (int i = 0; i < count; i++) {
            auto ev = epoll_.get_event(i);

            if (ev->events & EPOLLERR || ev->events & EPOLLHUP || ev->events & EPOLLRDHUP) {
                OnErrorEvent(errno);
                continue;
            }

            if (!is_connected_) {
                last_sendtime_ = GetCurrentMilliseconds();
                last_recvtime_ = last_sendtime_;
                OnConnected();
            }

            if (ev->events & EPOLLIN) {
                OnRecvEvent();
            }
            if (ev->events & EPOLLOUT) {
                OnSendEvent();
            }
        }
    }
}

void TcpClient::OnErrorEvent(int err) {
    log_info("TcpClient error event, err:%d", err);

    close_socket(fd_);

    auto obj = new InternalClientEventNode;
    if (is_connected_) {
        is_connected_ = false;
        obj->event = NetEvent::Closed;
    } else {
        obj->event = NetEvent::ConnectFailed;
    }
    obj->err_code = err;
    mpscq_.push(&obj->node);
    count_.FetchAdd(1, MemoryOrder::RELAXED);
    cv_.Signal();
}

void TcpClient::OnConnected() {
    is_connected_ = true;
    auto obj = new InternalClientEventNode;
    obj->event = NetEvent::ConnectSuccess;
    obj->err_code = 0;
    mpscq_.push(&obj->node);
    count_.FetchAdd(1, MemoryOrder::RELAXED);
    cv_.Signal();
}

void TcpClient::OnSendEvent() {
    AutoMutex g(&smtx_);

    if (!snd_buffer_.Empty() && SendImpl() != 0) {
        OnErrorEvent(errno);
        return;
    }

    if (snd_buffer_.Empty()) {
        epoll_.modify(fd_, reinterpret_cast<void *>(fd_), EPOLLIN | EPOLLET);
    } else {
        epoll_.modify(fd_, reinterpret_cast<void *>(fd_), EPOLLIN | EPOLLOUT | EPOLLET);
    }
}

void TcpClient::OnRecvEvent() {
    if (RecvImpl() != 0) {
        OnErrorEvent(errno);
        return;
    }

    if (snd_buffer_.Empty()) {
        epoll_.modify(fd_, reinterpret_cast<void *>(fd_), EPOLLIN | EPOLLET);
    } else {
        epoll_.modify(fd_, reinterpret_cast<void *>(fd_), EPOLLIN | EPOLLOUT | EPOLLET);
    }
}

bool TcpClient::SendMsg(const Slice &s) {
    if (shutdown_ || fd_ == -1 || s.empty()) {
        return false;
    }
    AutoMutex g(&smtx_);
    snd_buffer_.AddSlice(s);
    epoll_.modify(fd_, reinterpret_cast<void *>(fd_), EPOLLIN | EPOLLOUT | EPOLLET);
    return true;
}

void TcpClient::Disconnect() { close_socket(fd_); }

int TcpClient::RecvImpl() {

    int unused_space = TEMP_BUFFER_SIZE;
    int recv_bytes = 0;

    do {

        recv_bytes = ::recv(fd_, buffer_, unused_space, 0);

        if (recv_bytes == 0) {
            return -1;
        }

        if (recv_bytes < 0) {
            if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN) {
                break;
            }
            return -1;
        }

        last_recvtime_ = GetCurrentMilliseconds();

        TransferData(Slice(buffer_, recv_bytes));

    } while (recv_bytes == unused_space);
    return 0;
}

int TcpClient::SendImpl() {

    size_t count = 0;
    do {

        Slice slice = snd_buffer_.Front();
        int slen = ::send(fd_, slice.begin(), slice.size(), 0);

        if (slen == 0) {
            break;
        }

        if (slen < 0) {
            if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN) {
                break;
            }
            return -1;
        }

        last_sendtime_ = GetCurrentMilliseconds();

        snd_buffer_.MoveHeader((size_t)slen);
        count = snd_buffer_.SliceCount();

    } while (count > 0);
    return 0;
}

void TcpClient::TransferData(const Slice &s) {
    auto obj = new InternalClientEventNode;
    obj->event = NetEvent::Null;
    obj->err_code = 0;
    obj->data = s;
    mpscq_.push(&obj->node);
    count_.FetchAdd(1, MemoryOrder::RELAXED);
    cv_.Signal();
}

void TcpClient::WorkThread(void *) {
    while (!shutdown_) {
        if (count_.Load() < 1) {
            AutoMutex g(&cmtx_);
            cv_.Wait(&cmtx_, 1000);
            continue;
        }
        auto n = mpscq_.pop();
        if (n != nullptr) {
            count_.FetchSub(1, MemoryOrder::RELAXED);
            auto msg = reinterpret_cast<InternalClientEventNode *>(n);
            Dispatch(msg);
            delete msg;
        }
    }
}
void TcpClient::Dispatch(InternalClientEventNode *msg) {
    if (!msg->data.empty()) {
        rcv_buffer_.AddSlice(msg->data);
        if (rcv_buffer_.GetBufferLength() >= next_package_size_) {
            ParsingProtocol();
        }
        return;
    }

    service_->OnNetEvent(msg->event);
}

Slice TcpClient::ReadSliceFromRecvBuffer(size_t c) {
    Slice ret;
    size_t count = rcv_buffer_.SliceCount();
    for (size_t i = 0; i < count && i < c; i++) {
        ret += rcv_buffer_[i];
    }
    return ret;
}

void TcpClient::ParsingProtocol() {

    size_t cache_size = rcv_buffer_.GetBufferLength();
    int package_counter = 0;

    if (cache_size >= next_package_size_ && next_package_size_ > 0) {
        Slice packet = rcv_buffer_.GetHeader(next_package_size_);
        service_->OnMessage(packet);
        rcv_buffer_.MoveHeader(next_package_size_);
        cache_size -= next_package_size_;
        next_package_size_ = 0;
    }

    while (cache_size > 0) {
        size_t read_count = 1;
        int pack_len = 0;
        Slice packet;
        do {
            packet = ReadSliceFromRecvBuffer(read_count);
            pack_len = service_->MessageLength((uint8_t *)packet.buffer(), packet.length());
            if (pack_len < 0) {
                // discard buffer
                rcv_buffer_.MoveHeader(packet.length());
                cache_size = rcv_buffer_.GetBufferLength();
                break;
            }

            // equal 0 means we need more data
            if (pack_len == 0) {
                // reach the end of rcv_buffer
                if (read_count >= rcv_buffer_.SliceCount()) {
                    goto done;
                }
                // read more data to check msg header
                read_count++;
                continue;
            }

            // set next_package_size_ for next time
            if (cache_size < static_cast<size_t>(pack_len)) {
                next_package_size_ = pack_len;
                goto done;
            }

            // We got a whole packet with length pack_len
            break;

        } while (true);

        if (pack_len < 0) continue;

        if (packet.size() < static_cast<size_t>(pack_len)) {
            packet = rcv_buffer_.GetHeader(pack_len);
        } else {
            size_t n = packet.size() - pack_len;
            packet.RemoveTail(n);
        }
        service_->OnMessage(packet);
        rcv_buffer_.MoveHeader(pack_len);

        cache_size = rcv_buffer_.GetBufferLength();
    }
done:
    (void)package_counter;
}

} // namespace raptor
